
%% Conditional on imputed networks (and estimated A_hat), impute outcomes y1_educ and y1_gender
 

% To population with output
PE_Chain_Out = struct('BGD_hat_b',[],'A_hat_b',[],'G_ij_b',[],'burnin',burnin,'b_gap',b_gap,'PE_sim_reps',PE_sim_reps, 'PE_index',[], ...
    'V_C_b',[],'l_M_hat_b',[],'beta_M_hat_b',[],'sigma2_M_hat_b',[], ...
    'alpha_educ_m1_b',[],'alpha_educ_m2_b',[],'alpha_educ_m3_b',[],'alpha_educ_m4_b',[], ...
    'alpha_gender_m1_b',[],'alpha_gender_m2_b',[],'alpha_gender_m3_b',[],'alpha_gender_m4_b',[], ...
    'y1_educ_m1_b',[],'y1_educ_m2_b',[],'y1_educ_m3_b',[],'y1_educ_m4_b',[], ...
    'y1_gender_m1_b',[],'y1_gender_m2_b',[],'y1_gender_m3_b',[],'y1_gender_m4_b',[], ...
    'Sigma2_educ_m1_b',[],'Sigma2_educ_m2_b',[],'Sigma2_educ_m3_b',[],'Sigma2_educ_m4_b',[], ...
	'Sigma2_gender_m1_b',[],'Sigma2_gender_m2_b',[],'Sigma2_gender_m3_b',[],'Sigma2_gender_m4_b',[]);



    
%% Set initial stuff
load(['Mats_T2C.mat'])


y1_imp = struct('y1_educ_m1',NetOut_T2C.y1_educ,'y1_educ_m2',NetOut_T2C.y1_educ,'y1_educ_m3',NetOut_T2C.y1_educ,'y1_educ_m4',NetOut_T2C.y1_educ, ...
    'y1_gender_m1',NetOut_T2C.y1_gender,'y1_gender_m2',NetOut_T2C.y1_gender,'y1_gender_m3',NetOut_T2C.y1_gender,'y1_gender_m4',NetOut_T2C.y1_gender);



Cons_T2C.y1_educ_MISS = MISS_T2C.y1_educ_MISS;
Cons_T2C.y1_gender_MISS = MISS_T2C.y1_gender_MISS;




for c=1:chains
    load(['Net_EM_',num2str(sim_tol),'_ChainNo_',num2str(c),'_0.mat']); 
    b_max = size(Chain_Out.BGD_hat_b,2)
 
    
    z = burnin;
    
    while z < b_max
        
        if Net_OUT==0
            W_ij = exp(Chain_Out.G_ij_b(:,z)) + exp(Cons_T2C.Flip_ij*Chain_Out.G_ij_b(:,z));
        else
            W_ij = exp(Chain_Out.G_ij_b(:,z));
        end
        
        W_ij = Cons_T2C.W_trans*W_ij;
        W_ij_rep = zeros(Cons_T2C.n,Cons_T2C.n);
        k = 0;
        for s=1:Cons_T2C.schools
            W_ij_rep(k+1:k+Cons_T2C.sch_size(s),k+1:k+Cons_T2C.sch_size(s)) = reshape(W_ij(k+1:k+Cons_T2C.sch_size(s)^2,1),Cons_T2C.sch_size(s),Cons_T2C.sch_size(s));
            k = k + Cons_T2C.sch_size(s);
        end
        W_ij_imp = bsxfun(@rdivide,W_ij_rep,sum(W_ij_rep,2));   
        
        A_hat = Chain_Out.A_hat_b(:,z);
        


        % Initial Estimates

        for w=1:PE_sim_reps
            
            
            w
            %% Estimation Step
            [PE_out_imp_new] = Fn_EstimatePE_Mar2022(A_hat, W_ij_imp, y1_imp, Cons_T2C); 
            
            
            %% Imputation Step
            [y1_imp_new] = Fn_ImputePE_Mar2022(PE_out_imp_new, A_hat, W_ij_imp, Cons_T2C, y1_imp);
            
            y1_imp = y1_imp_new;
            PE_out_imp = PE_out_imp_new;

        end
        
        
        % Populate resutls into arrays
        PE_Chain_Out.PE_index = [PE_Chain_Out.PE_index, [c; z]];
        
        PE_Chain_Out.G_ij_b = [PE_Chain_Out.G_ij_b, Chain_Out.G_ij_b(:,z)];
        PE_Chain_Out.V_C_b = [PE_Chain_Out.V_C_b, Chain_Out.V_C_b(:,1)];
        PE_Chain_Out.A_hat_b = [PE_Chain_Out.A_hat_b, Chain_Out.A_hat_b(:,z)];
        PE_Chain_Out.BGD_hat_b = [PE_Chain_Out.BGD_hat_b, Chain_Out.BGD_hat_b(:,z)];        
        PE_Chain_Out.l_M_hat_b = [PE_Chain_Out.l_M_hat_b, Chain_Out.l_M_hat_b(:,z)]; 
        PE_Chain_Out.beta_M_hat_b = [PE_Chain_Out.beta_M_hat_b, Chain_Out.beta_M_hat_b(:,z)]; 
        PE_Chain_Out.sigma2_M_hat_b = [PE_Chain_Out.sigma2_M_hat_b, Chain_Out.sigma2_M_hat_b(:,z)]; 
        
        PE_Chain_Out.alpha_educ_m1_b = [PE_Chain_Out.alpha_educ_m1_b, PE_out_imp.alpha_educ_m1];
        PE_Chain_Out.alpha_educ_m2_b = [PE_Chain_Out.alpha_educ_m2_b, PE_out_imp.alpha_educ_m2];
        PE_Chain_Out.alpha_educ_m3_b = [PE_Chain_Out.alpha_educ_m3_b, PE_out_imp.alpha_educ_m3];
        PE_Chain_Out.alpha_educ_m4_b = [PE_Chain_Out.alpha_educ_m4_b, PE_out_imp.alpha_educ_m4];
        PE_Chain_Out.alpha_gender_m1_b = [PE_Chain_Out.alpha_gender_m1_b, PE_out_imp.alpha_gender_m1];
        PE_Chain_Out.alpha_gender_m2_b = [PE_Chain_Out.alpha_gender_m2_b, PE_out_imp.alpha_gender_m2];
        PE_Chain_Out.alpha_gender_m3_b = [PE_Chain_Out.alpha_gender_m3_b, PE_out_imp.alpha_gender_m3];
        PE_Chain_Out.alpha_gender_m4_b = [PE_Chain_Out.alpha_gender_m4_b, PE_out_imp.alpha_gender_m4];
        
        PE_Chain_Out.y1_educ_m1_b = [PE_Chain_Out.y1_educ_m1_b, y1_imp.y1_educ_m1];
        PE_Chain_Out.y1_educ_m2_b = [PE_Chain_Out.y1_educ_m2_b, y1_imp.y1_educ_m2];
        PE_Chain_Out.y1_educ_m3_b = [PE_Chain_Out.y1_educ_m3_b, y1_imp.y1_educ_m3];
        PE_Chain_Out.y1_educ_m4_b = [PE_Chain_Out.y1_educ_m4_b, y1_imp.y1_educ_m4];
        PE_Chain_Out.y1_gender_m1_b = [PE_Chain_Out.y1_gender_m1_b, y1_imp.y1_gender_m1];
        PE_Chain_Out.y1_gender_m2_b = [PE_Chain_Out.y1_gender_m2_b, y1_imp.y1_gender_m2];
        PE_Chain_Out.y1_gender_m3_b = [PE_Chain_Out.y1_gender_m3_b, y1_imp.y1_gender_m3];
        PE_Chain_Out.y1_gender_m4_b = [PE_Chain_Out.y1_gender_m4_b, y1_imp.y1_gender_m4];

        PE_Chain_Out.Sigma2_educ_m1_b = [PE_Chain_Out.Sigma2_educ_m1_b, PE_out_imp.Sigma2_educ_m1];
        PE_Chain_Out.Sigma2_educ_m2_b = [PE_Chain_Out.Sigma2_educ_m2_b, PE_out_imp.Sigma2_educ_m2];
        PE_Chain_Out.Sigma2_educ_m3_b = [PE_Chain_Out.Sigma2_educ_m3_b, PE_out_imp.Sigma2_educ_m3];
        PE_Chain_Out.Sigma2_educ_m4_b = [PE_Chain_Out.Sigma2_educ_m4_b, PE_out_imp.Sigma2_educ_m4];
        PE_Chain_Out.Sigma2_gender_m1_b = [PE_Chain_Out.Sigma2_gender_m1_b, PE_out_imp.Sigma2_gender_m1];
        PE_Chain_Out.Sigma2_gender_m2_b = [PE_Chain_Out.Sigma2_gender_m2_b, PE_out_imp.Sigma2_gender_m2];
        PE_Chain_Out.Sigma2_gender_m3_b = [PE_Chain_Out.Sigma2_gender_m3_b, PE_out_imp.Sigma2_gender_m3];
        PE_Chain_Out.Sigma2_gender_m4_b = [PE_Chain_Out.Sigma2_gender_m4_b, PE_out_imp.Sigma2_gender_m4];



        % Change z
        z = z + b_gap
    end
end

if Net_OUT==0
    save(['PE_Estimates_',num2str(sim_tol),'_OR.mat'],'PE_Chain_Out');
else
    save(['PE_Estimates_',num2str(sim_tol),'_OUT.mat'],'PE_Chain_Out');
end
    
